package first.group.ssh.packet;

import first.group.ssh.DataConsumer;
import first.group.ssh.readBuffer.ReadBuffer;
import src.AESCipherCTR;

import java.io.EOFException;
import java.io.IOException;
import java.nio.BufferOverflowException;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.channels.ReadableByteChannel;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.logging.Logger;
import javax.crypto.Cipher;
import javax.crypto.Mac;

public class PacketFactory implements DataConsumer {

    private final ByteBuffer buffer = ByteBuffer.allocate(65536);
    private final ByteBuffer decryptedBuffer = ByteBuffer.allocate(65536);
    private int decryptedLength = 0;
    private AESCipherCTR cipher = null;
    private Mac mac = null;
    private int seq = 0;

    public void fillIn(ReadBuffer buffer) {
        buffer.read(this.buffer);
    }

    public void setCipher(AESCipherCTR cipher, Mac mac) {
        this.cipher = cipher;
        this.mac = mac;
        if (decryptedBuffer.position() != 0) {
            throw new IllegalStateException("Changing cipher while decrypted buffer is not null");
        }
    }

    private Packet parseType(int type, ByteBuffer buffer) throws IOException {
        var parser = PacketParser.getParserFor(type);
        if (parser == null) {
            // TODO send SSH_MSG_UNIMPLEMENTED
            throw new IOException("Unknown type " + type);
        }
        return parser.parse(buffer);
    }

    private Packet parsePacketUnencrypted(ByteBuffer buffer) throws IOException {
        if (!PacketParser.bufferRemains(buffer, 8)) {
            return null;
        }
        var packetLen = buffer.getInt();
        if (!PacketParser.bufferRemains(buffer, packetLen)) {
            return null;
        }
        var paddingLen = Byte.toUnsignedInt(buffer.get());
        var type = Byte.toUnsignedInt(buffer.get());

        var origLimit = buffer.limit();
        buffer.limit(packetLen - paddingLen + 4);

        var packet = parseType(type, buffer);
        if (packet == null) {
            buffer.limit(origLimit);
            return null;
        }

        buffer.limit(origLimit);
        buffer.position(packetLen + 4);
        return packet;
    }

    private Packet parsePacketEncrypted() throws IOException {
        byte[] buf = new byte[cipher.getBlockSize()];
        byte[] macBuf = null;
        if (mac != null) {
            macBuf = new byte[mac.getMacLength()];
        }
        while (PacketParser.bufferRemains(buffer, cipher.getBlockSize())) {
            buffer.get(buf);

            if (decryptedLength == 0 || decryptedBuffer.position() < decryptedLength) {
                var decrypted = cipher.update(buf, 0, buf.length);
                assert decrypted != null;
                decryptedBuffer.put(decrypted);
            } else {
                decryptedBuffer.put(buf);
            }

            if (decryptedLength == 0) {
                // set decrypetedLength
                decryptedLength = decryptedBuffer.getInt(0) + 4;
                assert decryptedLength > 0;
                assert decryptedLength >= cipher.getBlockSize();
            } else {
                var macLength = 0;
                if (mac != null) {
                    macLength = mac.getMacLength();
                }
                var totalLength = decryptedLength + macLength;
                assert decryptedBuffer.position() <= totalLength;
                if (decryptedBuffer.position() == totalLength) {
                    decryptedBuffer.flip();
                    decryptedBuffer.limit(decryptedLength);
                    if (mac != null) {
                        var seqBuffer = ByteBuffer.allocate(4);
                        seqBuffer.order(ByteOrder.BIG_ENDIAN);
                        seqBuffer.putInt(seq);
                        seqBuffer.flip();
                        mac.update(seqBuffer);
                        mac.update(decryptedBuffer);
                    }

                    decryptedBuffer.position(0);
                    decryptedBuffer.limit(totalLength);
                    var packet = parsePacketUnencrypted(decryptedBuffer);
                    if (packet == null) {
                        throw new IOException("Parse packet failed");
                    }

                    if (mac != null) {
                        byte[] macData = mac.doFinal();
                        decryptedBuffer.get(macBuf);

                        if (!Arrays.equals(macBuf, macData)) {
                            throw new IOException("Mac different for seq " + seq + ": want " + Arrays.toString(macData) + ", but get " + Arrays.toString(macBuf));
                        }
                        mac.reset();
                    }

                    decryptedLength = 0;
                    decryptedBuffer.clear();

                    return packet;
                } else if (decryptedBuffer.position() >= totalLength) {
                    throw new AssertionError();
                }
            }
        }
        return null;
    }

    public Packet getPacket() throws IOException {
        if (cipher == null) {
            var total = buffer.position();
            buffer.position(0);
            buffer.limit(total);
            var packet = parsePacketUnencrypted(buffer);
            if (packet == null) {
                buffer.limit(buffer.capacity());
                buffer.position(total);
                return null;
            }
            var size = buffer.position();
            for (int i = 0; i < total - size; i++) {
                buffer.put(i, buffer.get(i + size));
            }
            buffer.position(total - size);
            buffer.limit(buffer.capacity());
            seq++;
            return packet;
        } else {
            var total = buffer.position();
            buffer.position(0);
            buffer.limit(total);
            var packet = parsePacketEncrypted();
            buffer.limit(buffer.capacity());

            var size = buffer.position();
            for (int i = 0; i < total - size; i++) {
                buffer.put(i, buffer.get(i + size));
            }
            buffer.position(total - size);
            if (packet != null) {
                seq++;
            }
            return packet;
        }
    }

    @Override
    public int fillData(ReadableByteChannel channel) throws IOException {
        var len = channel.read(buffer);
        if (len < 0) {
            throw new EOFException();
        }
        if (!buffer.hasRemaining()) {
            throw new IOException("Buffer overflow: Packet size exceeds " + buffer.capacity() + " bytes");
        }
        return len;
    }

    @Override
    public int fillData(byte[] data, int off, int len) throws IOException {
        try {
            buffer.put(data, off, len);
        } catch (BufferOverflowException ex) {
            throw new IOException("Buffer overflow: Packet size exceed.");
        }
        return len;
    }
}
